import matplotlib.pyplot as plt
import networkx as nx
from typing import Optional, List
import sys

"""基础的数据结构(Base Data Structure)
"""


class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

    def append(self, y) -> 'ListNode':
        if type(y) is type(self.val):
            self.next = ListNode(y)
        elif type(y) is ListNode:
            self.next = y
        return self.next

    def __str__(self):
        return str(self.val)


class ListConvert:
    @staticmethod
    def list_to_link(nums: list[Optional[int]]) -> Optional[ListNode]:
        """将数组转换为链表节点"""
        if not nums:
            return None

        head = ListNode(nums[0])
        cur = head
        for i in range(1, len(nums)):
            cur = cur.append(nums[i])

        return head

    @staticmethod
    def link_to_list(head: Optional[ListNode]):
        if not head:
            return []
        res = []
        curr = head
        while curr:
            res.append(curr.val)
            curr = curr.next
        return res


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

    def __str__(self):
        return str(self.val)


class TreeConverter:
    @staticmethod
    def list_to_tree(nums: list[Optional[int]]) -> Optional[TreeNode]:
        """将数组转换为二叉树

        Args:
            nums: 层序遍历的数组表示，如 [3,9,20,None,None,15,7]

        Returns:
            TreeNode: 二叉树的根节点
        """
        if not nums:
            return None

        root = TreeNode(nums[0])
        queue = [root]
        i = 1
        while queue and i < len(nums):
            node = queue.pop(0)
            # 左子节点
            if i < len(nums) and nums[i] is not None:
                node.left = TreeNode(nums[i])
                queue.append(node.left)
            i += 1
            # 右子节点
            if i < len(nums) and nums[i] is not None:
                node.right = TreeNode(nums[i])
                queue.append(node.right)
            i += 1
        return root

    @staticmethod
    def tree_to_list(root: Optional[TreeNode]) -> list[Optional[int]]:
        """将二叉树转换为数组

        Args:
            root: 二叉树的根节点

        Returns:
            List[Optional[int]]: 层序遍历的数组表示
        """
        if not root:
            return []

        result = []
        queue = [root]
        while queue:
            node = queue.pop(0)
            if node:
                result.append(node.val)
                queue.append(node.left)
                queue.append(node.right)
            else:
                result.append(None)

        # 移除末尾的None
        while result and result[-1] is None:
            result.pop()

        return result

    @staticmethod
    def draw_tree_from_root(root: Optional[TreeNode], figsize: tuple = (10, 6)) -> None:
        """使用networkx绘制二叉树

        Args:
            root: 树的根节点
            figsize: 图像大小
        """
        if not root:
            print("空树，无法绘制")
            return

        G = nx.DiGraph()
        pos = {}
        labels = {}

        def add_edges(node, x, y, level):
            if not node:
                return

            # 添加节点
            G.add_node(id(node))  # 使用节点的内存地址作为唯一标识
            pos[id(node)] = (x, y)
            labels[id(node)] = str(node.val)

            # 左子节点
            if node.left:
                G.add_edge(id(node), id(node.left))
                add_edges(node.left, x - 2 ** (-level), y - 1, level + 1)

            # 右子节点
            if node.right:
                G.add_edge(id(node), id(node.right))
                add_edges(node.right, x + 2 ** (-level), y - 1, level + 1)

        add_edges(root, 0, 0, 1)

        # 绘制图
        plt.figure(figsize=figsize)
        nx.draw(G, pos,
                with_labels=True,
                labels=labels,
                node_size=1000,
                node_color='lightblue',
                font_size=15,
                font_weight='bold',
                arrows=True,
                edge_color='gray',
                arrowsize=20)
        plt.title("Binary Tree Visualization")
        plt.axis('off')
        plt.show()

    @staticmethod
    def draw_tree_from_array(arr: list[Optional[int]], figsize: tuple = (10, 6)) -> None:
        """使用networkx通过数组绘制二叉树

        Args:
            arr: 层序遍历的数组表示
            figsize: 图像大小
        """
        if not arr:
            print("Empty Array")
            return

        root = TreeConverter.list_to_tree(arr)
        TreeConverter.draw_tree_from_root(root, figsize)


class AVLTree:
    def __init__(self):
        """初始化AVL树"""
        self.root = None

    def height(self, node: Optional[TreeNode]) -> int:
        """获取节点高度"""
        if not node:
            return 0
        return max(self.height(node.left), self.height(node.right)) + 1

    def balance_factor(self, node: Optional[TreeNode]) -> int:
        """计算平衡因子"""
        if not node:
            return 0
        return self.height(node.left) - self.height(node.right)

    def right_rotate(self, y: TreeNode) -> TreeNode:
        """右旋操作
        
             y                      x
            / \                   /   \
           x   T3   ======>     T1    y
          / \                        / \
         T1  T2                    T2  T3
        """
        x = y.left
        T2 = x.right

        x.right = y
        y.left = T2

        return x

    def left_rotate(self, x: TreeNode) -> TreeNode:
        """左旋操作
        
            x                         y
           / \                      /   \
          T1  y       ======>     x     T3
             / \                 / \
            T2  T3             T1  T2
        """
        y = x.right
        T2 = y.left

        y.left = x
        x.right = T2

        return y

    def insert(self, val: int) -> None:
        """插入节点"""

        def _insert(node: Optional[TreeNode], val: int) -> TreeNode:
            # 1. 执行标准BST插入
            if not node:
                return TreeNode(val)

            if val < node.val:
                node.left = _insert(node.left, val)
            elif val > node.val:
                node.right = _insert(node.right, val)
            else:  # 值相等，不做处理
                return node

            # 2. 更新平衡因子并进行平衡调整
            balance = self.balance_factor(node)

            # 左左情况
            if balance > 1 and val < node.left.val:
                return self.right_rotate(node)

            # 右右情况
            if balance < -1 and val > node.right.val:
                return self.left_rotate(node)

            # 左右情况
            if balance > 1 and val > node.left.val:
                node.left = self.left_rotate(node.left)
                return self.right_rotate(node)

            # 右左情况
            if balance < -1 and val < node.right.val:
                node.right = self.right_rotate(node.right)
                return self.left_rotate(node)

            return node

        self.root = _insert(self.root, val)

    def delete(self, val: int) -> None:
        """删除节点"""

        def get_min_value_node(node: TreeNode) -> TreeNode:
            current = node
            while current.left:
                current = current.left
            return current

        def _delete(node: Optional[TreeNode], val: int) -> Optional[TreeNode]:
            # 1. 执行标准BST删除
            if not node:
                return None

            if val < node.val:
                node.left = _delete(node.left, val)
            elif val > node.val:
                node.right = _delete(node.right, val)
            else:
                # 节点有0或1个子节点
                if not node.left:
                    return node.right
                elif not node.right:
                    return node.left

                # 节点有2个子节点
                # 获取中序后继节点（右子树最小值）
                successor = get_min_value_node(node.right)
                node.val = successor.val
                # 删除后继节点
                node.right = _delete(node.right, successor.val)

            # 如果树只有一个节点，直接返回
            if not node:
                return None

            # 2. 更新平衡因子并进行平衡调整
            balance = self.balance_factor(node)

            # 左左情况
            if balance > 1 and self.balance_factor(node.left) >= 0:
                return self.right_rotate(node)

            # 左右情况
            if balance > 1 and self.balance_factor(node.left) < 0:
                node.left = self.left_rotate(node.left)
                return self.right_rotate(node)

            # 右右情况
            if balance < -1 and self.balance_factor(node.right) <= 0:
                return self.left_rotate(node)

            # 右左情况
            if balance < -1 and self.balance_factor(node.right) > 0:
                node.right = self.right_rotate(node.right)
                return self.left_rotate(node)

            return node

        self.root = _delete(self.root, val)

    def to_list(self) -> List[Optional[int]]:
        """将AVL树转换为数组表示"""
        return TreeConverter.tree_to_list(self.root)

    def draw(self) -> None:
        """绘制AVL树"""
        TreeConverter.draw_tree_from_root(self.root)


def testAVL():
    avl = AVLTree()

    # 测试插入
    values = [10, 20, 30, 40, 50, 25]
    print("插入顺序:", values)
    for val in values:
        avl.insert(val)
        print(f"\n插入 {val} 后的树:")
        avl.draw()

    # 测试删除
    delete_values = [30, 20]
    print("\n删除操作:")
    for val in delete_values:
        avl.delete(val)
        print(f"\n删除 {val} 后的树:")
        avl.draw()


def total_size(obj, seen=None):
    """
    :param obj: 对象
    :param seen: 用来记录已经访问过的对象, 避免递归导致的无限循环
    :return:
    """
    if seen is None:
        seen = set()
    obj_id = id(obj)
    if obj_id in seen:
        return 0
    seen.add(obj_id)
    size = sys.getsizeof(obj)

    if isinstance(obj, dict):
        size += sum(total_size(k, seen) + total_size(v, seen) for k, v in obj.items())
    elif isinstance(obj, (list, tuple, set)):
        size += sum(total_size(i, seen) for i in obj)

    return size


# 测试代码
if __name__ == "__main__":
    head = ListNode(1)
    head.append(2).append(3)
    print(ListConvert.link_to_list(head))
